# seaborn已删除，只使用matplotlib进行可视化
import os
import warnings

# 只保留必要的导入
import matplotlib.pyplot as plt
import pandas as pd
from sentence_transformers import SentenceTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix, \
    classification_report
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'Heiti TC', 'SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

model = None
model_path = os.getenv('MODEL_PATH')


def train():
    """
    训练
    :return:
    """
    global model
    # --- 1. 创建示例情感数据集 ---
    print("创建示例情感数据集...")

    sample_data = {
        '文本': [
            # 正面情感 (label=1)
            "这个产品真的太棒了，我非常满意！",
            "服务态度很好，工作人员很热情",
            "质量超出了我的预期，强烈推荐",
            "今天心情特别好，阳光明媚",
            "这家餐厅的菜品味道很赞",
            "同事们都很友善，工作环境很棒",
            "这次旅行非常愉快，风景如画",
            "孩子们玩得很开心，活动组织得很好",
            "学到了很多新知识，收获满满",
            "朋友的生日聚会办得很成功",
            "非常感谢大家的帮助和支持",
            "这次合作非常成功，期待下次",
            "团队协作效果很好，目标达成了",
            "新功能使用起来很方便",
            "客户反馈非常积极正面",

            # 负面情感 (label=0)
            "这个产品质量太差了，完全不值这个价钱",
            "服务态度恶劣，让人很不舒服",
            "等了很久都没有得到回复，很失望",
            "今天遇到了很多麻烦事，心情糟糕",
            "这家店的食物难吃，环境也不好",
            "工作压力太大，每天都很疲惫",
            "天气太热了，让人烦躁不安",
            "交通堵塞严重，浪费了很多时间",
            "考试成绩不理想，感到很沮丧",
            "设备出现故障，影响了正常工作",
            "系统响应速度太慢，需要优化",
            "这个错误一直解决不了，很头疼",
            "预算不足，项目可能要延期",
            "沟通效率很低，浪费时间",
            "技术方案存在严重缺陷"
        ],
        'label': [
            # 正面=1 (15个)
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            # 负面=0 (15个)
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
        ]
    }

    df = pd.DataFrame(sample_data)
    print(f"数据集大小: {len(df)}")
    print(f"标签分布:\n{df['label'].value_counts()}")
    print("标签含义: 0=负面, 1=正面")

    # --- 2. 加载BGE-m3模型并提取特征 ---
    print("\n加载BGE-m3模型...")
    # MODEL_NAME = "BAAI/bge-m3"
    # model = SentenceTransformer(MODEL_NAME)
    embedding_model = SentenceTransformer(model_path)

    print("提取文本嵌入特征...")
    text_embeddings = embedding_model.encode(df['文本'].tolist(), show_progress_bar=True)
    print(f"嵌入特征形状: {text_embeddings.shape}")

    X = text_embeddings
    y = df['label']

    X, y = shuffle(X, y, random_state=42)
    print("数据已打乱。")

    # 将数据分成训练集和测试集
    # stratify=y 分层抽样
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

    print(f"训练集大小: {len(X_train)}")
    print(f"测试集大小: {len(X_test)}")
    print(f"训练集标签分布:\n{y_train.value_counts()}")
    print(f"测试集标签分布:\n{y_test.value_counts()}")

    # 使用训练集训练逻辑回归模型
    model = LogisticRegression(random_state=42, max_iter=1000)
    model.fit(X_train, y_train)
    print("\n模型训练完成！")

    # 在测试集上进行预测
    y_pred = model.predict(X_test)

    # 评估模型在测试集上的性能
    print("\n模型在测试集上的评估指标:")
    print(f"准确率 (Accuracy): {accuracy_score(y_test, y_pred):.4f}")
    print(f"精确率 (Precision): {precision_score(y_test, y_pred):.4f}")
    print(f"召回率 (Recall): {recall_score(y_test, y_pred):.4f}")
    print(f"F1-分数 (F1-Score): {f1_score(y_test, y_pred):.4f}")

    print("\n混淆矩阵 (Confusion Matrix):")
    print(confusion_matrix(y_test, y_pred))

    print("\n分类报告 (Classification Report):")
    print(classification_report(y_test, y_pred))


def predict(text: str):
    """
    预测
    :param text:
    :return:
    """
    global model
    embedding_model = SentenceTransformer(model_path)
    text_embedding = embedding_model.encode([text])
    predict_value = model.predict(text_embedding)[0]
    return predict_value


if __name__ == '__main__':
    # 训练
    train()

    # 预测
    text1 = '烦心事太多'
    predict_value = predict(text1)
    print(f'输入:{text1}')
    print(f'预测值:{predict_value}')
